import flask, json
from tokencost import calculate_prompt_cost

api = flask.Flask(__name__)

@api.route('/tokencost', methods=['post'])
def loginjosn():
    input = flask.request.json.get('input')
    output=flask.request.json.get('output')
    model = flask.request.json.get('model')

    input_cost = calculate_prompt_cost(input, model)
    output_cost=calculate_prompt_cost(output, model)

    ren = {
        "msg": "成功",
        "code": 200,
        "prompt_cost": input_cost+output_cost
    }

    return ren


if __name__ == '__main__':
    api.run(port=9000,host='0.0.0.0')  # 启动服务

    # debug=True,改了代码后，不用重启，它会自动重启

    # 'host='127.0.0.1'别IP访问地址
